import networkx as nx
from typing import Dict, Tuple, List, Set
from collections import defaultdict
import seaborn as sns

import self
from matplotlib import pyplot as plt

# from topo2 import build_hybrid_topology_shared_tor
from flow1 import TrafficGenerator
from topo1 import build_hybrid_topology_shared_tor


def allocate_switches_with_es(
        G: nx.Graph,
        T: List[Tuple[str, str]],
        D: Dict[Tuple[str, str], float],
        tor_switches: Set[str],
        optical_switches: Set[str],
        k_paths: int = 18
) -> Tuple[Dict, Dict]:
    """
    基于流量需求的ToR交换机和光交换机路径分配算法

    参数：
    G: 网络拓扑图（节点需标记类型，边需包含capacity属性）
    T: 通信对列表 [(t1, t2), ...]
    D: 流量需求字典 {(t1, t2): demand}
    tor_switches: ToR交换机集合
    optical_switches: 光交换机集合
    k_paths: 最大路径查找数量

    返回：
    S: 交换机分配方案 {(t1, t2): (tor1, tor2)}
    P: 路径分配方案 {(t1, t2): path}
    """

    # 初始化数据结构
    S = {} # 交换机分配方案
    P = {} # 路径分配方案

    # 初始化每条链路剩余容量（使用有序边元组）
    remaining_capacity = {}
    for u, v in G.edges():
        edge = tuple(sorted((u, v)))
        remaining_capacity[edge] = G[u][v]['capacity']

    print("==============================================")
    print("列出初始化的每条链路剩余容量")
    for edge, capacity in remaining_capacity.items():
        print(f"{edge}: {capacity}")

    # 初始化光交换机占用状态（字典类型）
    # switch_occupied = {node: False for node in optical_switches}
    # print("初始交换机的占用状态：",switch_occupied)

    # 按流量需求降序排序通信对
    sorted_pairs = sorted(T, key=lambda x: D[x], reverse=True)
    print("降序排序后的通信对",sorted_pairs)

    # 处理每个通信对(每个通信对就是ToR通信对）
    for t1, t2 in sorted_pairs:
        demand = D[(t1, t2)]
        print(f"{t1,t2}的通信对需求为：",demand)
        # print("输出每个通信节点对的信息")
        # print(t1,t2)

        # 获取连接的ToR交换机
        # print("ToR交换机的集合：",tor_switches)
        # print(f"{t1}的邻接ToR交换机:",G.neighbors(t1))
        # print(f"{t2}的邻接ToR交换机:",G.neighbors(t2))
        # tor1 = [n for n in G.neighbors(t1) if n in tor_switches]
        # tor2 = [n for n in G.neighbors(t2) if n in tor_switches]

        # if not tor1 or not tor2:
        #     print(f"节点{t1}没有邻接ToR交换机")
        #     S[(t1, t2)] = None
        #     P[(t1, t2)] = None
        #     continue
        # if not tor1:
        #     print(f"节点 {t1} 没有邻接 ToR 交换机")
        #     S[(t1, t2)] = None
        #     P[(t1, t2)] = None
        #     continue
        # elif not tor2:
        #     print(f"节点 {t2} 没有邻接 ToR 交换机")
        #     S[(t1, t2)] = None
        #     P[(t1, t2)] = None
        #     continue

        # 查找所有可行路径
        all_paths = []
        try:
            # 使用Yen's算法查找前k条路径
            paths = nx.shortest_simple_paths(G, t1, t2, weight='hops')
            all_paths.extend([(t1, t2, p) for p, _ in zip(paths, range(k_paths))])
        except nx.NetworkXNoPath:
            continue

        # 使用 all_simple_paths 函数查找从 t1 到 t2 的所有简单路径
        # try:
        #     paths = nx.all_simple_paths(G, t1, t2)
        #     for path in paths:
        #         print(f"从 {t1} 到 {t2} 的路径: {path}")
        #     all_paths.extend([(t1, t2, p) for p in paths])
        # except nx.NetworkXNoPath:
        #     print(f"从 {t1} 到 {t2} 没有找到路径")
        #     continue

        # 输出所有路径
        # for path in all_paths:
        #     print(f"从 {path[0]} 到 {path[1]} 的路径: {path[2]}")


        # for src in t1:
        #     for dst in t2:
        #         try:
        #             # 使用Yen's算法查找前k条路径
        #             paths = nx.shortest_simple_paths(G, src, dst, weight='hops')
        #             all_paths.extend([(src, dst, p) for p, _ in zip(paths, range(k_paths))])
        #         except nx.NetworkXNoPath:
        #             continue

        # 过滤可行路径
        feasible_paths = []
        for src, dst, path in all_paths:
            if is_path_feasible(path, optical_switches, remaining_capacity, demand):
                feasible_paths.append((t1, t2, path))

        if not feasible_paths:
            print("没有可行的路径")
            S[(t1, t2)] = None
            P[(t1, t2)] = None
            continue

        # 选择最优路径
        best_path = select_optimal_path(feasible_paths, remaining_capacity)
        t1, t2, path = best_path

        # 记录分配方案
        S[(t1, t2)] = (t1, t2)
        P[(t1, t2)] = path

        # 更新网络资源
        update_resources(path, remaining_capacity, demand, optical_switches)

    return S, P


def is_path_feasible(
        path: List[str],
        optical_switches: Set[str],
        # switch_occupied: Dict[str, bool],
        remaining_capacity: Dict[Tuple[str, str], float],
        demand: float
) -> bool:
    """检查路径是否可行"""
    # 检查光交换机占用状态
    # for node in path:
    #     if node in optical_switches and switch_occupied[node]:
    #         print(f"由于当前光交换机{node}已被占用，无法分配该路径")
    #         return False

    # 检查链路容量
    for i in range(len(path) - 1):
        u, v = sorted((path[i], path[i + 1]))
        edge = (u, v)
        if remaining_capacity[edge] < demand:
            # print(f"可行性检查：此时的链路{edge}容量还剩{remaining_capacity[edge]}由于链路容量不足，不满足当前的流量需求{demand}无法分配该路径")
            return False

    return True


def select_optimal_path(
        feasible_paths: List[Tuple[str, str, List[str]]],
        remaining_capacity: Dict[Tuple[str, str], float]
) -> Tuple[str, str, List[str]]:
    """选择最优路径（最大瓶颈容量）"""
    max_capacity = -1 # 表示设置最大瓶颈容量为正数
    # best_path = None # 初始化最优路径为None
    best_path = []
    for src, dst, path in feasible_paths:
        min_cap = float('inf')
        for i in range(len(path) - 1):
            u, v = sorted((path[i], path[i + 1]))
            min_cap = min(min_cap, remaining_capacity[(u, v)])

        if min_cap > max_capacity:
            max_capacity = min_cap
            best_path = (src, dst, path)

    return best_path


def update_resources(
        path: List[str],
        remaining_capacity: Dict[Tuple[str, str], float],
        demand: float,
        optical_switches: Set[str],
        # switch_occupied: Dict[str, bool]
) -> None:
    """更新网络资源状态"""
    # 更新链路容量
    for i in range(len(path) - 1):
        u, v = sorted((path[i], path[i + 1]))
        remaining_capacity[(u, v)] -= demand
        # print(f"{u,v}被分配的流量为{demand},更新之后的链路容量为：",remaining_capacity[(u, v)])

    # 标记光交换机
    # for node in path:
    #     if node in optical_switches:
    #         switch_occupied[node] = True
    #         print(f"更新光交换机{node}已经被占用")


# 测试用例
if __name__ == "__main__":
    # 创建测试拓扑
    G = build_hybrid_topology_shared_tor()  # 从函数中获取拓扑
    # print("---------------------------------------------------------")
    # 检查图的节点
    print("图的节点G.nodes:", G.nodes()) # 获取到的是所有的节点，包括交换机和服务器
    generator = TrafficGenerator(num_nodes=4) # 从函数中获取流量生成器，输入的16是ToR的数量
    print("===========生成流量矩阵D_dict=====================")
    D = generator.generate_all_to_all() # 生成所有节点对之间的流量（ToR节点对之间的流量）


    # 将 numpy.ndarray 转换为字典
    D_dict = {(i, j): D[i][j] for i in range(D.shape[0]) for j in range(D.shape[1])}
    print("D_dict",D_dict)
    # # 获取 ToR 交换机和光交换机集合
    tor_switches = {node for node, data in G.nodes(data=True) if data.get('type') == 'tor'}
    print("tor_switches",tor_switches)
    optical_switches = {node for node, data in G.nodes(data=True) if data.get('type') == 'optical_switch'}
    print("光交换机节点为：",optical_switches)
    print("--------------------------------------------------------------------------------")
    print("D_dict.items()",D_dict.items())
    # 生成通信对列表
    print("===========生成流量矩阵D_dict=====================")

    T = []
    for pair, demand in D_dict.items():
        # print("pair:",pair)
        # print("demand:",demand)
        if demand > 0:
            T.append(pair)
    for pair in T:
        print("pair:",pair) # 输出的通信对的格式为（1,2），表示ToR1和ToR2

    # 将键从 (i, j) 转换为 ("ToR{i}", "ToR{j}")
    D_str_keys = {(f"ToR{i+1}", f"ToR{j+1}"): value for (i, j), value in D_dict.items()}
    # print("D_str_keys:",D_str_keys)

    # 将值从 numpy.ndarray 转换为浮点数
    D_float_values = {k: float(v) for k, v in D_str_keys.items()}

    # 生成通信对列表
    T = [pair for pair, demand in D_float_values.items() if demand > 0]
    print("生成通信对列表",T)

    # 检查通信对节点是否在图中
    for pair in T:
        for node in pair:
            if node not in G.nodes():
                print(f"节点 {node} 不在图 G 中")

    # 调用分配算法
    S, P = allocate_switches_with_es(G, T, D_float_values, tor_switches, optical_switches)

    # 输出结果
    print("交换机分配方案:")
    for pair, assignment in S.items():
        print(f"{pair} -> {assignment}")

    print("\n路径分配方案:")
    for pair, path in P.items():
        print(f"{pair} -> {path}")


    # 根据流量矩阵绘制热力图
    plt.figure(figsize=(10, 8))  # 设置图形的大小
    sns.heatmap(D, annot=True, cmap='YlGnBu', fmt='.2f', cbar=True)  # 绘制热力图
    plt.title('Flow Demand Heatmap')  # 设置标题
    plt.xlabel('ToR Nodes')  # 设置x轴标签
    plt.ylabel('ToR Nodes')  # 设置y轴标签
    plt.show()  # 显示热力图
